package com.wzy.aspect;

import cn.hutool.extra.servlet.ServletUtil;
import com.wzy.annotation.LimitRequest;
import com.wzy.domain.ResponseResult;
import com.wzy.service.UserCommonService;
import com.wzy.util.RedisCache;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.AfterReturning;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import java.util.concurrent.TimeUnit;

import static com.wzy.constant.RedisConstants.LIMIT_REQUEST_KEY;
import static com.wzy.constant.SystemConstants.DEFINED;
import static com.wzy.constant.SystemConstants.TRUE;
import static com.wzy.enums.AppHttpCodeEnum.LIMIT_ERROR;

/**
 * 接口请求限制切面类
 */
@Aspect
@Component
@Slf4j
@RequiredArgsConstructor
public class LimitRequestAspect
{
    private final HttpServletRequest request;
    private final UserCommonService userCommonService;
    private final RedisCache redisCache;

    // 定义切点
    @Pointcut("@annotation(limitRequest)")
    public void includeService(LimitRequest limitRequest)
    {
    }

    /**
     * 拦截被限流请求
     */
    @Around(value = "includeService(limitRequest)", argNames = "pjp,limitRequest")
    public Object doAround(ProceedingJoinPoint pjp, LimitRequest limitRequest) throws Throwable
    {
        String key = getKey(limitRequest);
        String redisKey = LIMIT_REQUEST_KEY + request.getRequestURI() + ":" + key;

        if (Boolean.FALSE.equals(redisCache.hasKey(redisKey)))
        {
            return pjp.proceed();
        }
        Integer uCount = redisCache.getCacheObject(redisKey);
        if (!limitRequest.limitAdmin() && userCommonService.isAdmin(request))
        {
            // 对管理员放行
            return pjp.proceed();
        }
        if (uCount >= limitRequest.count())
        {
            // 超过次数，不执行目标方法
            if (limitRequest.tip().equals(DEFINED))
            {
                log.info("被拦截的请求 => URL:[{}] IP/USER:[{}]", request.getRequestURI(), key);
                return ResponseResult.errorResult(LIMIT_ERROR, limitRequest.description());
            }
            else if (limitRequest.tip().equals(TRUE))
            {
                log.info("被拦截的请求 => URL:[{}] IP/USER:[{}]", request.getRequestURI(), key);
                return ResponseResult.errorResult(LIMIT_ERROR, redisCache.getExpireTime(redisKey) + "秒 后再试");
            }
            else
            {
                return ResponseResult.okResult();
            }
        }
        return pjp.proceed();
    }

    /**
     * 计数器（只记录成功的请求）
     */
    @AfterReturning(value = "includeService(limitRequest)", returning = "result", argNames = "jp,limitRequest,result")
    public void doAfterReturning(JoinPoint jp, LimitRequest limitRequest, ResponseResult result)
    {
        String key = getKey(limitRequest);
        String redisKey = LIMIT_REQUEST_KEY + request.getRequestURI() + ":" + key;

        if (limitRequest.limitAdmin() || !userCommonService.isAdmin(request))
        {
            if (Boolean.FALSE.equals(redisCache.hasKey(redisKey)))
            {
                redisCache.setCacheObject(redisKey, 0, limitRequest.time(), TimeUnit.MILLISECONDS);
            }
            else
            {
                redisCache.increment(redisKey);
            }
        }
    }

    public String getKey(LimitRequest limitRequest)
    {
        String key;

        if (limitRequest.type().equals("IP"))
        {
            // 获取真实IP
            key = ServletUtil.getClientIP(request);
        }
        else
        {
            key = userCommonService.getRedisUserId(request).toString();
        }
        return key;
    }

}